import os
from flask import Flask, render_template, request, jsonify, send_from_directory, url_for
from models.yolo_model import YOLOModel
from models.yolo1_model import YOLOModel1
from utils.Text_to_Image import text_to_image
from jinja2 import TemplateNotFound
app = Flask(__name__)
model = YOLOModel()
model1 = YOLOModel1()

# 配置文件夹
app.config['RESULT_FOLDER'] = os.path.join(app.root_path, 'result')
os.makedirs(app.config['RESULT_FOLDER'], exist_ok=True)

@app.route('/')
def index():
    return render_template('main.html')  # 默认加载 main.html

@app.route('/<page>')
def load_page(page):
    if page != 'get_next_frame':
        return render_template(f"{page}")
    else:
        return url_for('get_next_frame')

@app.route('/图像分割')
def image_segmentation():
    return render_template('image_segmentation.html')

@app.route('/目标检测', methods=['POST', 'GET'])
def object_detection():
    return render_template('object_detection.html')

@app.route('/图像生成')
def image_generation():
    return render_template('image_generation.html')

@app.route('/detect', methods=['POST'])
def detect():
    file = request.files.get('file')
    is_video = request.form.get('is_video', 'false') == 'true'

    if not file:
        return jsonify({"error": "No file provided"}), 400

    input_path = os.path.join("data", file.filename)
    output_path = os.path.join("static", "results", file.filename)
    os.makedirs(os.path.dirname(input_path), exist_ok=True)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    file.save(input_path)

    if is_video:
        global processed_frames
        processed_frames = []

    is_model = request.form.get('model')
    if is_model == "model1":
        segmentation_stats = model1.segment(input_path, output_path, is_video=is_video)
        return jsonify({"stats": segmentation_stats, "output_path": output_path})
    else:
        detection_stats = model.detect(input_path, output_path, is_video=is_video)
        return jsonify({"stats": detection_stats, "output_path": output_path})

@app.route('/get_next_frame', methods=['POST'])
def get_next_frame():
    is_model = request.form.get('model')
    if is_model == "model1":
        frame_path, finished = model1.get_next_frame()
    else:
        frame_path, finished = model.get_next_frame()

    if frame_path is None:
        return jsonify({"finished": True}), 200

    if not os.path.exists(frame_path):
        return jsonify({"error": "Frame not found"}), 404

    if finished:
        os.remove("./static/result")
        os.makedirs("./static/result", exist_ok=True)

    return jsonify({"frame_path": frame_path, "finished": finished})

@app.route('/generate', methods=['POST'])
def generate():
    data = request.get_json()
    desc = data.get('description')
    if not desc:
        return jsonify({"error": "描述不能为空"}), 400

    image_path = text_to_image(desc)
    static_path = os.path.join('static', 'generated_images')
    os.makedirs(static_path, exist_ok=True)
    output_path = os.path.join(static_path, os.path.basename(image_path))
    os.rename(image_path, output_path)

    image_url = f"/static/generated_images/{os.path.basename(output_path)}"
    return jsonify({"image_url": image_url}), 200

@app.route('/result/<filename>')
def get_image(filename):
    return send_from_directory(app.config['RESULT_FOLDER'], filename)

if __name__ == '__main__':
    app.run(debug=True,host="0.0.0.0", port=5000)
